#include "StdAfx.h"
#include "ObjectTracker.h"
#include "StackWalker.h"

#include <sstream>

ObjectTracker::~ObjectTracker(void)
{
}

void ObjectTracker::Initialize(CComPtr<ICorProfilerInfo2> pCorProfilerInfo)
{
	m_TypeTracker.RegisterForClassLoad(HandleClassLoad, this);

	//OutputDebugStringA("ObjectTracker Initialized");
	m_pCorProfilerInfo = pCorProfilerInfo;
}

STDMETHODIMP ObjectTracker::SurvivingReferences(ULONG cSurvivingObjectIDRanges,
												ObjectID objectIDRangeStart[  ],
												ULONG cObjectIDRangeLength[  ])
{
	std::vector<ObjectID> garbageCollegedObjectIds;

	for (OBJECT_CREATIONS_ITERATOR iter = m_resolvedObjectCreations.begin();
		iter != m_resolvedObjectCreations.end();
		++iter)
	{
		ObjectID id = iter->first;

		bool objectSurvivedGC = false;
		for (ULONG i = 0; i<cSurvivingObjectIDRanges; ++i)
		{
			// If id lies within the contiguous range
			if (id >= objectIDRangeStart[i] && id < (objectIDRangeStart[i] + cObjectIDRangeLength[i]))
			{
				objectSurvivedGC = true;
				break;
			}
		}

		if (!objectSurvivedGC)
			garbageCollegedObjectIds.push_back(id);
	}

	for (std::vector<ObjectID>::const_iterator iter = garbageCollegedObjectIds.begin();
		iter != garbageCollegedObjectIds.end();
		++iter)
	{
		m_resolvedObjectCreations.erase(*iter);
	}

	return S_OK;
}

STDMETHODIMP ObjectTracker::MovedReferences( ULONG cMovedObjectIDRanges,
											ObjectID oldObjectIDRangeStart[  ],
											ObjectID newObjectIDRangeStart[  ],
											ULONG cObjectIDRangeLength[  ])
{
	std::vector<std::pair<ObjectID, ObjectCreationSnapshot> > currentObjectIds(m_resolvedObjectCreations.begin(),
		m_resolvedObjectCreations.end());

	for (std::vector<std::pair<ObjectID, ObjectCreationSnapshot> >::iterator iter = currentObjectIds.begin();
		iter != currentObjectIds.end();
		++iter)
	{
		ObjectID id = iter->first;

		bool objectSurvivedGC = false;

		ULONG rangeIndex = 0;
		ULONG offsetInOldRange = 0;
		for (ULONG i = 0; i<cMovedObjectIDRanges; ++i)
		{
			// If id lies within the contiguous range
			if (id >= oldObjectIDRangeStart[i] && id < (oldObjectIDRangeStart[i] + cObjectIDRangeLength[i]))
			{
				rangeIndex = i;
				offsetInOldRange = id - oldObjectIDRangeStart[i];
				break;
			}
		}

		m_resolvedObjectCreations.erase(id);

		ULONG calculatedId = newObjectIDRangeStart[rangeIndex] + offsetInOldRange;
		if (calculatedId < (newObjectIDRangeStart[rangeIndex] + cObjectIDRangeLength[rangeIndex]))
		{
			m_resolvedObjectCreations[calculatedId] = iter->second;
		}
	}

	return S_OK;
}

STDMETHODIMP ObjectTracker::ObjectAllocated(ObjectID objectId, ClassID classId)
{
	//OutputDebugStringA("ObjectAllocated Start");
	CLRClass clrClass;
	HRESULT hr = m_TypeTracker.GetCLRClassFor(classId, &clrClass);

	if (FAILED(hr) || PassesFilter(clrClass.fullTypeName))
	{	
		ObjectCreationSnapshot snapshot;
		snapshot.classId = classId;
		snapshot.objectId = objectId;

		StackWalker stackWalker(m_TypeTracker);

		hr = stackWalker.WalkStack(m_pCorProfilerInfo);

		if (SUCCEEDED(hr))
			snapshot.stack = stackWalker.GetStack();

		if (FAILED(hr))
		{
			m_unResolvedObjectCreations[classId] = snapshot;
		}
		else
		{
			m_resolvedObjectCreations[objectId] = snapshot;
		}
	}
	//OutputDebugStringA("ObjectAllocated End");
	return S_OK;
}

HRESULT ObjectTracker::GetCreationSnapshot(ObjectID objectId, ObjectCreationSnapshot *creationSnapshot) const
{
	OBJECT_CREATIONS_ITERATOR iter = m_resolvedObjectCreations.find(objectId);

	if (iter == m_resolvedObjectCreations.end())
		return E_FAIL;

	ObjectCreationSnapshot &snapshot = const_cast<ObjectCreationSnapshot &>(iter->second);

	CLRClass clrClass;
	if (SUCCEEDED(m_TypeTracker.GetCLRClassFor(snapshot.classId, &clrClass)))
		snapshot.clrClass = clrClass;

	*creationSnapshot = snapshot;

	return S_OK;
}

void ObjectTracker::SetTypeFilter(const std::vector<std::wstring> &typeFilters)
{
	for (std::vector<std::wstring>::const_iterator iter = typeFilters.begin(); iter != typeFilters.end(); ++iter)
	{
		OutputDebugStringW((*iter).c_str());
		m_TypeFilters.insert(*iter);
	}
}

void ObjectTracker::HandleClassLoad(const CLRClass &classLoad, void *data)
{
	ObjectTracker *pTracker = static_cast<ObjectTracker *>(data);

	UNRESOLVED_OBJECT_CREATIONS_ITERATOR iter = pTracker->m_unResolvedObjectCreations.find(classLoad.classId);

	if (iter != pTracker->m_unResolvedObjectCreations.end())
	{
		const ObjectCreationSnapshot &snapshot = iter->second;
		pTracker->m_unResolvedObjectCreations.erase(classLoad.classId);
		pTracker->m_resolvedObjectCreations[snapshot.objectId] = snapshot;
	}
}

bool ObjectTracker::PassesFilter(const std::wstring &typeName)
{
	return m_TypeFilters.find(typeName) != m_TypeFilters.end();
}